"""
车牌检测模块
使用YOLOv5进行车牌定位
"""

import torch
import numpy as np
import cv2
from typing import List, Tuple, Dict, Any, Optional
from pathlib import Path
from ultralytics import YOLO
import requests
import os


class PlateDetector:
    """
    车牌检测器类
    
    功能：
    - 使用YOLOv5检测图像中的车牌
    - 支持批量检测
    - 自动下载预训练模型
    """
    
    def __init__(self, config: Optional[dict] = None):
        """
        初始化车牌检测器
        
        参数:
            config (dict, optional): 配置字典
        """
        self.config = config or {}
        self.model_path = self.config.get('model_path', 'models/yolov5_plate.pt')
        self.conf_threshold = self.config.get('conf_threshold', 0.5)
        self.nms_threshold = self.config.get('nms_threshold', 0.45)
        self.img_size = self.config.get('img_size', 640)
        
        # 设置设备
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        
        # 加载模型
        self.model = self._load_model()
        
    def _load_model(self) -> YOLO:
        """
        加载YOLOv5模型
        
        返回:
            YOLO: 加载的模型
        """
        model_path = Path(self.model_path)
        
        # 如果模型不存在，尝试下载或使用默认模型
        if not model_path.exists():
            print(f"模型文件 {model_path} 不存在")
            
            # 确保模型目录存在
            model_path.parent.mkdir(parents=True, exist_ok=True)
            
            # 尝试下载预训练的车牌检测模型
            if self._download_model():
                print("模型下载成功")
            else:
                # 使用YOLOv5默认模型
                print("使用默认YOLOv5模型")
                model_path = 'yolov5nu.pt'  # 使用兼容的模型
        
        # 加载模型
        try:
            model = YOLO(str(model_path))
            model.to(self.device)
            print(f"模型加载成功，使用设备: {self.device}")
            return model
        except Exception as e:
            print(f"模型加载失败: {e}")
            # 使用默认的兼容模型
            try:
                model = YOLO('yolov5nu.pt')
                model.to(self.device)
                print("使用默认兼容模型 yolov5nu.pt")
                return model
            except Exception as e2:
                print(f"默认模型加载也失败: {e2}")
                raise e2
            
    def _download_model(self) -> bool:
        """
        下载预训练的车牌检测模型
        
        返回:
            bool: 是否下载成功
        """
        try:
            # 使用ModelLoader下载模型
            from ..models import ModelLoader
            from ..utils import Config
            
            # 加载配置
            config = Config()
            loader = ModelLoader(
                models_dir=str(Path(self.model_path).parent),
                config=config.get_all()
            )
            
            # 下载yolov5_plate模型
            model_path = loader.load_model('yolov5_plate')
            if model_path:
                # 如果下载成功，将模型文件重命名为期望的路径
                if model_path != Path(self.model_path):
                    import shutil
                    shutil.move(str(model_path), self.model_path)
                return True
                
        except Exception as e:
            print(f"使用ModelLoader下载失败: {e}")
            
        # 如果ModelLoader失败，回退到直接下载
        model_urls = [
            "https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov5nu.pt",
            "https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n.pt"
        ]
        
        for url in model_urls:
            try:
                print(f"正在从 {url} 下载模型...")
                response = requests.get(url, stream=True, timeout=30)
                response.raise_for_status()
                
                # 保存模型文件
                with open(self.model_path, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        f.write(chunk)
                        
                return True
            except Exception as e:
                print(f"从 {url} 下载失败: {e}")
                continue
                
        return False
        
    def detect(self, image: np.ndarray) -> List[Dict[str, Any]]:
        """
        检测图像中的车牌
        
        参数:
            image (np.ndarray): 输入图像 (BGR格式)
            
        返回:
            List[Dict[str, Any]]: 检测结果列表，每个结果包含:
                - bbox: 边界框坐标 [x1, y1, x2, y2]
                - confidence: 置信度
                - class_name: 类别名称
        """
        # 运行检测
        results = self.model(image, 
                           conf=self.conf_threshold,
                           iou=self.nms_threshold,
                           imgsz=self.img_size)
        
        # 解析结果
        detections = []
        for r in results:
            if r.boxes is not None:
                for box in r.boxes:
                    # 获取边界框坐标
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    
                    # 获取置信度
                    confidence = float(box.conf[0])
                    
                    # 获取类别
                    class_id = int(box.cls[0])
                    class_name = self.model.names.get(class_id, 'plate')
                    
                    detection = {
                        'bbox': [int(x1), int(y1), int(x2), int(y2)],
                        'confidence': confidence,
                        'class_name': class_name
                    }
                    detections.append(detection)
                    
        return detections
        
    def detect_batch(self, images: List[np.ndarray]) -> List[List[Dict[str, Any]]]:
        """
        批量检测图像中的车牌
        
        参数:
            images (List[np.ndarray]): 图像列表
            
        返回:
            List[List[Dict[str, Any]]]: 每张图像的检测结果
        """
        # 批量运行检测
        results = self.model(images, 
                           conf=self.conf_threshold,
                           iou=self.nms_threshold,
                           imgsz=self.img_size)
        
        # 解析每张图像的结果
        batch_detections = []
        for r in results:
            detections = []
            if r.boxes is not None:
                for box in r.boxes:
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    confidence = float(box.conf[0])
                    class_id = int(box.cls[0])
                    class_name = self.model.names.get(class_id, 'plate')
                    
                    detection = {
                        'bbox': [int(x1), int(y1), int(x2), int(y2)],
                        'confidence': confidence,
                        'class_name': class_name
                    }
                    detections.append(detection)
                    
            batch_detections.append(detections)
            
        return batch_detections
        
    def visualize_detections(self, image: np.ndarray, 
                           detections: List[Dict[str, Any]],
                           color: Tuple[int, int, int] = (0, 255, 0),
                           thickness: int = 2) -> np.ndarray:
        """
        在图像上可视化检测结果
        
        参数:
            image (np.ndarray): 原始图像
            detections (List[Dict]): 检测结果
            color (Tuple[int, int, int]): 边界框颜色
            thickness (int): 线条粗细
            
        返回:
            np.ndarray: 绘制了检测结果的图像
        """
        result_image = image.copy()
        
        for det in detections:
            bbox = det['bbox']
            confidence = det['confidence']
            
            # 绘制边界框
            cv2.rectangle(result_image, 
                         (bbox[0], bbox[1]), 
                         (bbox[2], bbox[3]), 
                         color, thickness)
            
            # 绘制标签
            label = f"Plate: {confidence:.2f}"
            label_size, _ = cv2.getTextSize(label, 
                                           cv2.FONT_HERSHEY_SIMPLEX, 
                                           0.5, 1)
            
            # 绘制标签背景
            cv2.rectangle(result_image,
                         (bbox[0], bbox[1] - label_size[1] - 4),
                         (bbox[0] + label_size[0], bbox[1]),
                         color, -1)
            
            # 绘制标签文本
            cv2.putText(result_image, label,
                       (bbox[0], bbox[1] - 2),
                       cv2.FONT_HERSHEY_SIMPLEX,
                       0.5, (255, 255, 255), 1)
                       
        return result_image 